from matplotlib import pyplot as plt
import numpy as np
from pathlib import Path
import serial
import cion.data as cion_data
from cion.gui import scan_lines_widgets, scan_generator_pop
from cion.sequencer import Sequence, plot_sequence
### 下面是cion调用硬件接口的部分
### 对于特定的实验环境，需要用户自主地实现这些硬件接口,以顺利调用相应的硬件设备
import cion.hardware.serial_comm as hardwarecomm # FPGA串口通信接口（如有必要）
import cion.hardware.awg_control as awgcomm # AWG以太网通信（如有必要）

all = []

# 确定使用模式：真实硬件模式（debug_mode == False）或测试模式（）
class Global_Setting:
    __debug_mode: bool = False
    #__debug_mode: bool = True
    __ser = None

    @classmethod
    def get_ser(cls):
        return cls.__ser

    @classmethod
    def get_debug_mode(cls):
        return cls.__debug_mode

    @classmethod
    def set_ser(cls, ser):
        __ser = ser

    @classmethod
    def set_debug_mode(cls, mode:bool):
        cls.__debug_mode = mode


def load_sequencies(file_path):
    seq_file = Path(file_path)
    if seq_file.is_file():
        execfile(seq_file)
    return

### An example to test importlib.reload()
class A:
    def __init__(self):
        self.aa = 1
        print('ab')


# 用于离子阱实验的主类，每次开始新的实验都应该创建一个Experiment实例。
# 例如：exp = Experiment(2, chapter_dict = FPGA_dict, port = 'COM3')
# 表示开始一次实验，此实验中离子数为2，FPGA对应的chapter表为FPGA_dict, 端口号为COM3
class Experiment:
    def __init__(self, ion_number = -1, rf_pick = 854, melt_status = False, eta = 0.098, chapter_dict = None, port = None, ip_address = None):
        self.ion_number = ion_number
        self.rf_pick = rf_pick
        self.mode_frequency = np.random.random(ion_number)
        self.melt_status = melt_status
        self.Lamb_Dicke_parameter = eta
        self.channel = [2,4,0,6]
        self.clock_period = 5E-3
        self.repeat = 100
        self.ext_trig = 1
        # Different version of FPGA could have different channel_return.
        # If FPGA doesn't work as expected, please check whether channel_return is correcly setted.
        self.channel_return = 24
        self.port = port
        self.awg_switch = False
        self.threshold = None
        self.state_flag = False

        self.__chapter_dict = None
        if chapter_dict != None:
            self.__chapter_dict = hardwarecomm.chapter_process(chapter_dict)
        else:
            print('Warning! You did not specify the chapter data for this experiments, the FPGA utility will be deactivated.')

        # 一次实验可以创建多个Sequence，全部存储在seq_list里面
        self.seq_list = []
        self.last_sequence = Sequence(ion_number=ion_number, chapter_dict=self.__chapter_dict)
        all.clear()
        [all.append(i) for i in range(ion_number)]

        assert(isinstance(ion_number, int) and ion_number >= 1)
        assert(isinstance(melt_status, bool))

        # debug模式下：
        ### 假装PC连接了底层各类硬件，并且将“实验文件”保存至当前目录下的test_datafile_dir文件夹里
        # 实验模式下：
        ### 结束上一次实验的串口连接，并且为本次实验新建串口通讯
        ### 连接awg的以太网地址
        if Global_Setting.get_debug_mode() == True:
            self.data = cion_data.DataDoc(disk_path=r'test_datafile_dir',debug_mode=Global_Setting.get_debug_mode())
            print("/********************\nCaution! You are using debug mode. \
This enables you to test some functions locally. But you cannot access the actual hardware in this mode.\n********************/")
        else:
            self.data = cion_data.DataDoc(debug_mode=Global_Setting.get_debug_mode())

            self.connect_hardware(port, ip_address)

##### 与硬件相关的接口函数，在硬件配置发生改变时，这些函数需要完全重构
    def connect_hardware(self, port = None, ip_address = None):
        #连接serial port
        if port != None:
            if Global_Setting.get_ser() != None:
                if Global_Setting.get_ser().isOpen():
                    Global_Setting.get_ser().close()
            this_ser = serial.Serial(
                port=port,
                baudrate=4000000,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_TWO,
                bytesize=serial.EIGHTBITS
            )
            Global_Setting.set_ser(this_ser)
            self.ser = this_ser
        # 连接awg
        if ip_address != None:
            self.awg_switch = True
            self.awg_interface = awgcomm.AWG_Control(ip_address)

    def reopen_serial_port(self):
        assert self.port != None
        if Global_Setting.get_ser() != None:
            if Global_Setting.get_ser().isOpen():
                Global_Setting.get_ser().close()
        this_ser = serial.Serial(
            port=self.port,
            baudrate=4000000,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_TWO,
            bytesize=serial.EIGHTBITS
        )
        Global_Setting.set_ser(this_ser)
        self.ser = this_ser
#####
#####

    # 在Experiment里创建一个新的Sequence实例
    def new_sequence(self):
        seq = Sequence(ion_number=self.ion_number, chapter_dict=self.__chapter_dict)
        self.seq_list.append(seq)
        self.last_sequence = seq
        return seq

    def check_sequence(self, sequence):
        if sequence == None:
            sequence = self.last_sequence
        return sequence

    def set_path_prefix(self, prefix):
        self.path_prefix = prefix

    def set_ion_number(self, ion_number):
        self.ion_number = ion_number
        all.clear()
        [all.append(i) for i in range(ion_number)]
        if ion_number > len(self.channel):
            print("Alert! The ion number in current setup is larger than the predefined channel numbers.\n Try using\
            member function \'set_channel\' to reset channel number.")

    def sweep_generator(self,n,sequence=None):
            i = 0
            x = np.array([0.5] * n)
            while True:
                counts = self.run_once(sequence=sequence)
                counts = np.array([sum(counts)/self.repeat]*n)
                i += 1
                yield i,counts

    def sweep(self, repeat=100, steps=100, fps=64, win_size=200, sequence=None):
        generator = self.sweep_generator(self.ion_number, sequence=sequence)
        reader, viewer, button = scan_lines_widgets(
            n=self.ion_number,
            hist_bins=10,
            producer=lambda: next(generator),
            win_size=win_size,
            sampling_rate=64)
        return

    def print_info_all():
        #print(global_parameters)
        #print(sequences)
        pass

    def print_info_global():
        print("debug_mode:", Global_Setting.get_debug_mode())
        print("serial_port:", Global_Setting.get_ser())
        pass

    def conditional_run_once(self, mode, repeat=100, ext_trig=1, sequence = None):
        if sequence == None:
            sequence = self.last_sequence
        # debug模式下不会访问硬件
        if Global_Setting.get_debug_mode() == True:
            print(sequence.perform_sequence_serial())
        # 实验模式下才会真正与硬件交互
        else:
            self.run_once(repeat=repeat, ext_trig=ext_trig)

    def time_scan_generator(self, scan_list, ranges, sequence = None):
        '''
        scan_list: a list of labels, like '['label1', 'rx1', 'rx2', ...]'
        ranges: duration list to be scanned
        (Old definition: ranges: a tuple, its length equivalent to scan_list's length)
        '''
        if sequence == None:
            sequence = self.last_sequence
        n = self.ion_number
        x = np.array([0.5] * n)

        #assert len(scan_list) == len(ranges)
        #Independently scan durations

        if isinstance(ranges[0],(np.ndarray,list)):
            scan_ranges = range(len(ranges[0]))
            def update_duration(i):
                for label, durations in zip(scan_list, ranges):
                    sequence.set_parameters(label=label, duration=durations[i], batch_mode=True)

                sequence.regenerate_sequence()
                if self.awg_switch == True:
                    self.send_awg_data(sequence)
                #Run experiment
                raw_counts = self.run_once()
                counts = np.array([sum(raw_counts)/self.repeat])
                self.data.append_raw_data(i, raw_counts)
                return counts[-1]
        #Globally scan durations
        else:
            scan_ranges = ranges
            def update_duration(duration):
                if isinstance(scan_list,(np.ndarray,list)):
                    for label in scan_list:
                        sequence.set_parameters(label=label, duration=duration, batch_mode=True)
                else:
                    assert isinstance(scan_list,str)
                    sequence.set_parameters(label=scan_list, duration=duration, batch_mode=True)
                ####更新硬件参数
                sequence.regenerate_sequence()
                if self.awg_switch == True:
                    self.send_awg_data(sequence)
                ####
                #Run experiment
                raw_counts = self.run_once()
                counts = np.array([sum(raw_counts)/self.repeat])
                self.data.append_raw_data(duration, counts)
                return counts[-1]
        # sequence updated
        # now please get data for the expermients
        #TODO
        return scan_ranges, update_duration

    def time_scan(self, scan_list, ranges, fps=64, sequence=None, cycle=1):
        #Update rawdata
        if sequence == None:
            sequence = self.last_sequence
        self.data.title = "TimeScan"
        self.data.reset_raw_data(xtype="Duration")

        scan_ranges, update_duration = self.time_scan_generator(scan_list, ranges, sequence)
        scan_generator_pop(func=update_duration, scan_list=scan_ranges, title="Duration (us)", cycle=cycle)
        self.data.generate_json()
        return

    def freq_scan_generator(self, scan_list, ranges, sequence):
        n = self.ion_number
        x = np.array([0.5] * n)
        #assert len(scan_list) == len(ranges)
        #Independently scan durations

        if isinstance(ranges[0],(np.ndarray,list)):
            scan_ranges=range(len(ranges[0]))
            def update_freq(i):
                for label, freqs in zip(scan_list, ranges):
                    sequence.set_parameters(label=label, freq=freqs[i], batch_mode=True)
                sequence.regenerate_sequence()
                sequence.update_hardwares(scan_list)
                if self.awg_switch == True:
                    self.send_awg_data(sequence)
                #Run experiment
                raw_counts = np.array(self.run_once())
                counts = np.array([sum(raw_counts)/self.repeat])
                self.data.append_raw_data(freqs[0], raw_counts)
                return counts[-1]
        #Globally scan durations
        else:
            scan_ranges=ranges
            def update_freq(freq):
                if isinstance(scan_list,(np.ndarray,list)):
                    for label in scan_list:
                        sequence.set_parameters(label=label, freq=freq, batch_mode=True)
                else:
                    assert isinstance(scan_list,str)
                    #print(scan_list)
                    #print(sequence.labelled_pulse)
                    sequence.set_parameters(label=scan_list, freq=freq, batch_mode=True)
                sequence.regenerate_sequence()
                ####更新硬件参数
                sequence.update_hardwares(scan_list)
                if self.awg_switch == True:
                    self.send_awg_data(sequence)
                ####

                #Run experiment
                raw_counts = self.run_once()
                counts = np.array([sum(raw_counts)/self.repeat]*n)
                self.data.append_raw_data(freq, raw_counts)
                return counts[-1]
        return scan_ranges, update_freq

    # 扫频率
    def freq_scan(self, scan_list, ranges, fps=64, sequence=None, cycle=1):
        if sequence == None:
            sequence = self.last_sequence
        self.data.title = "FreqScan"
        self.data.reset_raw_data(xtype="Frequency")
        scan_ranges, update_freq = self.freq_scan_generator(scan_list, ranges, sequence)
        scan_generator_pop(func=update_freq, scan_list=scan_ranges, title="Frequency (MHz)", cycle=cycle)
        self.data.generate_json()
        sequence.reset_hardwares(scan_list)
        return

    def amp_scan_generator(self, scan_list, ranges, sequence):
        n = self.ion_number
        x = np.array([0.5] * n)
        #assert len(scan_list) == len(ranges)
        #Independently scan durations

        if isinstance(ranges[0],(np.ndarray,list)):
            scan_ranges = range(len(ranges[0]))
            def update_amp(i):
                for label, amps in zip(scan_list, ranges):
                    sequence.set_parameters(label=label, amp=amps[i], batch_mode=True)

                sequence.update_hardwares(scan_list)
                sequence.regenerate_sequence()
                if self.awg_switch == True:
                    self.send_awg_data(sequence)
                #Run experiment
                raw_counts = self.run_once()
                counts = np.array([sum(raw_counts)/self.repeat]*n)
                self.data.append_raw_data(i, raw_counts)
                return counts[-1]
        #Globally scan durations
        else:
            scan_ranges = ranges
            def update_amp(amp):
                if isinstance(scan_list,(np.ndarray,list)):
                    for label in scan_list:
                        sequence.set_parameters(label=label, amp=amp, batch_mode=True)
                else:
                    assert isinstance(scan_list,str)
                    sequence.set_parameters(label=scan_list, amp=amp, batch_mode=True)

                sequence.update_hardwares(scan_list)
                sequence.regenerate_sequence()
                if self.awg_switch == True:
                    self.send_awg_data(sequence)
                #Run experiment
                raw_counts = self.run_once()
                counts = np.array([sum(raw_counts)/self.repeat]*n)
                self.data.append_raw_data(amp, raw_counts)
                return counts[-1]
        file_name = self.data.generate_json()
        return scan_ranges, update_amp

    # 扫幅度
    def amp_scan(self, scan_list, ranges, fps=64, sequence=None, cycle=1):
        if sequence == None:
            sequence = self.last_sequence
        self.data.title = "AmpScan"
        self.data.reset_raw_data(xtype="Amp")

        scan_ranges, update_amp = self.amp_scan_generator(scan_list, ranges, sequence)
        scan_generator_pop(func=update_amp, scan_list=scan_ranges, title="Amp", cycle=cycle)
        self.data.generate_json()
        return

    # 画出sequence的时序图
    def print_sequence(self, size=(10, 4), sequence=None):
        if sequence == None:
            sequence = self.last_sequence
        #figure = plt.figure(figsize=size)
        plot_sequence(sequence.gate_sequence, remap_time=False, given_size=size)
        plt.tight_layout()
        plt.show()
        pass

    def save_info_to_database():
        pass

##### hardware interface functions
##### 与硬件相关的接口函数，在硬件配置发生改变时，这些函数需要完全重构

    # 硬件相关：更新本次实验所用的FPGA chapter并对其进行格式处理
    def update_chapter_dict(self, chapter_dict):
        self.__chapter_dict = hardwarecomm.chapter_process(chapter_dict)

    # 硬件相关：设定离子的读取通道（实验的数据通道有可能多于离子数，此时需要确定每一个离子分别从哪个通道读取）
    def set_channel(self, channel):
        '''
        reset channel and convert matrix
        '''
        assert(self.ion_number == len(channel))
        self.channel = channel

        m = 2*self.ion_number-1
        self.convert_matrix = np.zeros((self.ion_number, m), dtype=int)
        for i in range(self.ion_number):
            self.convert_matrix[i][channel[i]] = 1


    # 通过awg_interface发送AWG波形数据到AWG硬件上
    def send_awg_data(self, sequence=None, print_flag=False):
        if sequence == None:
            sequence = self.last_sequence
        sequence.perform_sequence_awg()
        if Global_Setting.get_debug_mode() == True:
            fig, axs= plt.subplots(nrows = 1, ncols = self.ion_number, figsize=(27, 6))
            i=0
            for ax in axs:
                ax.plot(sequence.awg_data[i])
                i += 1
            plt.show()
        else:
            # 保证有awg通道打开
            if self.awg_switch == True:
                self.awg_interface.send_data(sequence.awg_data, print_flag=print_flag)
            else:
                print("Caution! You are trying to send awg data while no awg interface defined. Check the ip address?")

    # 实际运行一个序列时只需要和FPGA进行串口通信，awg的数据是提前传输好的
    def run_once(self, ext_trig=1, sequence=None):
        if sequence == None:
            sequence = self.last_sequence
        if Global_Setting.get_debug_mode() == False:
            ### 如果处于实验模式下，将编码后的数据通过串口发送给FPGA，并读取FPGA的返回结果至result
            repeat = self.repeat
            threshold = [self.threshold,self.state_flag]
            pulses = sequence.perform_sequence_serial()
            pulses = hardwarecomm.pre_binary(pulses, self.clock_period)
            packets = [hardwarecomm.packets_generator(pulses, ext_trig, repeat), int(repeat*self.channel_return)]
            ### 如果处于debug模式下，则是直接生成相应长度的随机序列，存入result中
            result = hardwarecomm.seq_send(self.ser, packets, threshold)
        # 如果处于debug模式下，则是直接生成相应长度的随机序列，存入result中
        else:
            result = np.random.rand(self.ion_number*self.repeat) * 20
            import time
            time.sleep(0.005)
        return result
#####
#####
